import time
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets,transforms
import torch.nn.functional as F

class MyNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, out_dim):
        super(MyNet, self).__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, out_dim)
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.sigmoid(self.layer2(x))
        return x

p = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train = datasets.FashionMNIST('fashionMINIST_data', train=True, transform=p, download=True)
test = datasets.FashionMNIST('fashionMNIST_data', train=False, transform=p, download=True)

data_train = DataLoader(train, batch_size=16, shuffle=True)
data_test = DataLoader(test, batch_size=16, shuffle=False)

# a=iter(train)
# img,label=next(a)
# print(img.shape)

net = MyNet(28*28*1, 300, 10)

loss_fun = nn.CrossEntropyLoss()
optimer = torch.optim.SGD(net.parameters(), lr = 0.01, momentum=0.9)

train_ls = []
train_acc = []
epoch_num = 15
start = time.time()
for epoch in range(epoch_num):
    all_loss = 0.0
    all_acc = 0.0
    for img, label in data_train:
        optimer.zero_grad()
        out = net.forward(img.view(img.shape[0], -1))
        loss = loss_fun(out, label)
        loss.backward()
        optimer.step()
        pred = out.argmax(dim=1)
        correct = (pred==label).sum().item()
        acc = correct/img.shape[0]
        all_loss += loss.item()
        all_acc += acc
    avg_loss = all_loss/len(data_train)
    avg_acc = all_acc/len(data_train)
    train_ls.append(avg_loss)
    train_acc.append(avg_acc)
    print(f'epoch:{epoch+1}, loss:{avg_loss:.4f}, acc:{avg_acc:.4f}')
end = time.time()

print(f'训练所用时间:{end-start}s')

plt.subplot(121)
plt.plot(range(epoch_num), train_ls)
plt.subplot(122)
plt.plot(range(epoch_num), train_acc)
plt.show()

test_loss = 0.0
test_acc = 0.0
with torch.no_grad():
    for img, label in data_test:
        out = net.forward(img.view(img.shape[0], -1))
        loss = loss_fun(out, label)
        pred = out.argmax(dim=1)
        correct =(pred==label).sum().item()
        acc = correct/img.shape[0]
        test_loss += loss.item()
        test_acc += acc
    print(f'loss:{test_loss/len(data_test):.4f}, acc:{test_acc/len(data_test):.4f}')

